import wx

import presets
from event_bus import EventBus
from exceptions import LocationNotSetException
from uis import ProjectInfo
from view_models import project_state
from wxutils import call_after


class ProjectSetup(ProjectInfo):

    def __init__(self, parent):
        super().__init__(parent)

        self.nn_presets.SetItems(list(presets.NetworkInputDimension.keys()))

        self.Bind(wx.EVT_BUTTON, self.save_project_info, self.save_project)
        self.Bind(wx.EVT_BUTTON, self.select_project_location, self.open_location)
        self.Bind(wx.EVT_BUTTON, self.upgrade_project_data, self.upgrade_project)
        self.Bind(wx.EVT_CHOICE, self.apply_network_template, self.nn_presets)
        self.Bind(wx.EVT_TEXT, self.change_size, self.sample_long_side)
        self.Bind(wx.EVT_TEXT, self.correspond_calculate_rates, self.training_rate)
        self.Bind(wx.EVT_TEXT, self.correspond_calculate_rates, self.validation_rate)

        EventBus().register_topic(topic="project_state_updated", handler=self.update_project_state)

    @call_after
    def update_project_state(self, _):
        self.project_name.ChangeValue(project_state.current.name)
        self.project_location.ChangeValue(project_state.current.location)
        self.export_category.ChangeValue(project_state.current.export_file)
        self.export_sample.ChangeValue(project_state.current.export_sample)
        self.export_validate.ChangeValue(project_state.current.export_validate)
        self.sample_type.SetSelection(self.sample_type.GetItems().index(project_state.current.sample_type))
        self.sample_long_side.ChangeValue(f"{project_state.current.sample_long_side}")
        self.training_groups.ChangeValue(f"{project_state.current.training_groups}")
        self.training_rate.ChangeValue(f"{project_state.current.training_rate * 100:.2f}")
        self.validation_rate.ChangeValue(f"{project_state.current.validation_rate * 100:.2f}")
        self.upgrade_project.Show(project_state.current.version < 2)

    def select_project_location(self, _):
        if not project_state.has_location():
            dir_select = wx.DirDialog(self, "选择项目建立位置")
            if dir_select.ShowModal() == wx.ID_OK:
                project_state.current.location = dir_select.GetPath()
                EventBus().post_topic(event=None, topic="project_state_updated")
        else:
            wx.MessageBox("项目文件已经在指定位置建立且不可移动", "不能完成操作")

    def apply_network_template(self, event: wx.CommandEvent):
        if event.GetString() in presets.NetworkInputDimension:
            preset = presets.NetworkInputDimension[event.GetString()]
            self.sample_long_side.ChangeValue(f"{preset}")

    def change_size(self, _):
        self.nn_presets.SetSelection(0)

    def correspond_calculate_rates(self, event: wx.CommandEvent):
        if event.GetEventObject() is self.training_rate:
            source = self.training_rate
            related = self.validation_rate
        elif event.GetEventObject() is self.validation_rate:
            source = self.validation_rate
            related = self.training_rate
        else:
            return
        source_value = float(event.GetString())
        if source_value > 100:
            source.ChangeValue("100.00")
            source_value = 100
        related.ChangeValue(f"{100 - source_value}")

    def save_project_info(self, _):
        try:
            project_state.current.name = self.project_name.GetValue()
            project_state.current.location = self.project_location.GetValue()
            project_state.current.export_file = self.export_category.GetValue()
            project_state.current.export_sample = self.export_sample.GetValue()
            project_state.current.export_validate = self.export_validate.GetValue()
            project_state.current.sample_type = self.sample_type.GetStringSelection()
            project_state.current.sample_long_side = int(self.sample_long_side.GetValue())
            project_state.current.training_groups = int(self.training_groups.GetValue())
            project_state.current.training_rate = float(self.training_rate.GetValue()) / 100
            project_state.current.validation_rate = float(self.validation_rate.GetValue()) / 100

            project_state.save()
            wx.MessageBox("项目保存成功。", "保存成功")
            EventBus().post_topic(event=None, topic="enable_functions")
            EventBus().post_topic(event=None, topic="project_state_updated")
        except LocationNotSetException:
            wx.MessageBox("项目位置尚未设定，请先设定项目位置。", "保存失败")

    def upgrade_project_data(self, _):
        if wx.MessageBox("是否重新计算全部样本源及样本的指纹？", "操作确认", wx.YES_NO, self) == wx.YES:
            EventBus().post_topic(event=True, topic="toggle_upgrade_fingerprint")
            from exports.base import pool
            pool.submit(self.upgrade_fingerprints)

    @staticmethod
    def upgrade_fingerprints():
        from PIL import Image
        from models import Sample, Source
        import imagehash
        from peewee import DatabaseError

        print("Upgrade source:")
        for s in Source.select():
            try:
                old_fp = s.fingerprint
                new_fp = str(imagehash.dhash(Image.open(s.path)))
                Source.update(fingerprint=new_fp).where(Source.path == s.path).execute()
                Sample.update(source_fingerprint=new_fp).where(Sample.source == s.path).execute()
                print(f"{old_fp} -> {new_fp}")
            except DatabaseError as e:
                print(e)

        print("Upgrade sample:")
        for s in Sample.select():
            try:
                old_fp = s.target_fingerprint
                source_image = Image.open(s.source)
                cropped_image = source_image.crop((max([0, s.crop_x]),
                                                   max([0, s.crop_y]),
                                                   min([s.crop_x + s.crop_width, source_image.width]),
                                                   min([s.crop_y + s.crop_height, source_image.height])))
                new_image = Image.new(mode="RGBA", size=(s.crop_width, s.crop_height))
                dest_x = abs(s.crop_x) if s.crop_x < 0 else 0
                dest_y = abs(s.crop_y) if s.crop_y < 0 else 0
                new_image.alpha_composite(cropped_image.convert("RGBA"), (dest_x, dest_y))
                new_fp = str(imagehash.dhash(new_image))
                Sample.update(target_fingerprint=new_fp).where(Sample.serial == s.serial).execute()
                print(f"{old_fp} -> {new_fp}")
            except DatabaseError as e:
                print(e)

        project_state.current.version = 2
        project_state.save()
        print("All done.")
        EventBus().post_topic(event=False, topic="toggle_upgrade_fingerprint")
        EventBus().post_topic(event=None, topic="project_state_updated")
